from vnpy_ctastrategy.backtesting import BacktestingEngine
from vnpy.trader.constant import Interval
from datetime import datetime
from vnpy.trader.constant import Exchange, Product

class CustomBacktestingEngine(BacktestingEngine):
    """支持本地CSV数据的回测引擎"""
    
    def __init__(self, data_adapter=None):
        super().__init__()
        self.data_adapter = data_adapter
    
    def load_data_from_csv(self, symbol: str, start: datetime, end: datetime):
        """
        从CSV文件加载数据到回测引擎
        """
        if not self.data_adapter:
            raise ValueError("需要先设置data_adapter")
        
        # 确定交易所
        if symbol.endswith('.SSE'):
            exchange = Exchange.SSE
            pure_symbol = symbol.replace('.SSE', '')
        elif symbol.endswith('.SZSE'):
            exchange = Exchange.SZSE
            pure_symbol = symbol.replace('.SZSE', '')
        else:
            exchange = Exchange.SSE  # 默认上交所
            pure_symbol = symbol
        
        # 加载数据
        bars = self.data_adapter.load_bar_data(
            symbol=pure_symbol,
            exchange=exchange,
            interval=Interval.DAILY,
            start=start,
            end=end
        )
        
        if not bars:
            print(f"未找到数据: {symbol}")
            return False
        
        # 清空现有数据
        self.history_data.clear()
        
        # 添加数据到引擎
        for bar in bars:
            self.history_data.append(bar)
        
        print(f"成功加载 {len(bars)} 条数据: {symbol}")
        return True
